"""
@description: Core training logic using Hugging Face Trainer, integrating model, data, and configurations.
@author: yangkw
@date: 2025-05-02
"""
# -*- coding: utf-8 -*-
import logging
import os
import torch
from typing import Optional # 导入 Optional
from transformers import (
    TrainingArguments as HfTrainingArguments, # 使用官方的 TrainingArguments
    Trainer,
    DataCollatorForLanguageModeling,
    set_seed,
    EarlyStoppingCallback # 导入早停回调
)
from peft import get_peft_model, LoraConfig, TaskType, prepare_model_for_kbit_training
import wandb
# 不再需要手动导入和初始化 Accelerator
# from accelerate import Accelerator

# 从同级目录导入我们定义的模块和类
from .utils import (
    ModelArguments,
    DataTrainingArguments,
    QuantizationArguments,
    PeftArguments,
    WandbArguments,
    logger # 使用 utils 中配置好的 logger
)
from .model_loader import load_model_and_tokenizer
from .data_processor import get_processed_datasets

def run_training(
    model_args: ModelArguments,
    data_args: DataTrainingArguments,
    training_args: HfTrainingArguments, # 接收官方的 TrainingArguments
    quant_args: QuantizationArguments,
    peft_args: PeftArguments,
    wandb_args: WandbArguments,
    overwrite_cache: bool = False, # 添加 overwrite_cache 参数
    early_stopping_callback: Optional[EarlyStoppingCallback] = None # 添加早停回调参数
):
    """
    执行模型微调的主要函数。

    Args:
        model_args: 模型相关参数。
        data_args: 数据相关参数。
        training_args: Hugging Face Trainer 的训练参数。
        quant_args: 量化相关参数。
        peft_args: PEFT/LoRA 相关参数。
        wandb_args: Wandb 相关参数。
    """
    # --- 0. 初始化和设置 ---
    logger.info("="*30 + " 开始训练流程 " + "="*30)
    logger.info(f"训练参数 (TrainingArguments): {training_args}")
    logger.info(f"模型参数 (ModelArguments): {model_args}")
    logger.info(f"数据参数 (DataTrainingArguments): {data_args}")
    logger.info(f"量化参数 (QuantizationArguments): {quant_args}")
    logger.info(f"PEFT 参数 (PeftArguments): {peft_args}")
    logger.info(f"Wandb 参数 (WandbArguments): {wandb_args}")
    logger.info(f"Overwrite Cache: {overwrite_cache}") # 记录传入的 overwrite_cache 值

    # 不再需要手动初始化 Accelerator，Trainer 会处理
    # accelerator = Accelerator(log_with=training_args.report_to)
    # logger.info(f"Accelerator 状态: {accelerator.state}")

    # 设置随机种子 (确保所有进程都设置)
    # Trainer 初始化时也会设置种子，但提前设置更保险
    set_seed(training_args.seed)
    logger.info(f"设置随机种子: {training_args.seed}")

    # 设置 Wandb (如果启用)
    # Trainer 会自动处理初始化，但我们可以在主进程提前设置环境变量或进行检查
    # 使用 training_args.process_index 判断主进程
    if training_args.process_index == 0 and "wandb" in training_args.report_to:
        if wandb_args.wandb_project:
            os.environ["WANDB_PROJECT"] = wandb_args.wandb_project
            logger.info(f"主进程设置 Wandb 项目: {wandb_args.wandb_project}")
        if wandb_args.wandb_entity:
            os.environ["WANDB_ENTITY"] = wandb_args.wandb_entity
            logger.info(f"主进程设置 Wandb 实体: {wandb_args.wandb_entity}")
        if wandb_args.wandb_name:
            os.environ["WANDB_NAME"] = wandb_args.wandb_name
            logger.info(f"主进程设置 Wandb 运行名称: {wandb_args.wandb_name}")
        # os.environ["WANDB_LOG_MODEL"] = "checkpoint" # 可以配置自动上传模型检查点

    # --- 1. 加载模型和 Tokenizer ---
    # Trainer 和 datasets 会在 accelerate 环境下自动处理主进程优先下载
    # 不再需要手动使用 accelerator.main_process_first()
    # with accelerator.main_process_first():
    model, tokenizer = load_model_and_tokenizer(model_args, quant_args)

    # --- 2. 加载和处理数据 ---
    # 同样，Trainer 和 datasets 会自动处理
    # with accelerator.main_process_first():
    processed_datasets = get_processed_datasets(
        tokenizer=tokenizer,
        data_args=data_args,
        training_args=training_args, # 传递 HfTrainingArguments
        overwrite_cache=overwrite_cache, # 传递 overwrite_cache
        num_proc=training_args.dataloader_num_workers # 使用 dataloader_num_workers 作为 map 的进程数
    )

    train_dataset = processed_datasets.get("train")
    eval_dataset = processed_datasets.get("validation")

    if train_dataset is None:
        raise ValueError("未能加载或处理训练数据集。")
    logger.info(f"训练数据集大小: {len(train_dataset)}")
    if eval_dataset:
        logger.info(f"验证数据集大小: {len(eval_dataset)}")

    # --- 3. 配置 PEFT (LoRA) ---
    if peft_args.use_peft:
        logger.info("启用 PEFT (LoRA) 微调。")

        # !! 关键改动：仅在启用 PEFT 时，才为 k-bit 训练准备模型 !!
        if quant_args.load_in_4bit or getattr(quant_args, 'load_in_8bit', False): # 兼容可能存在的 8bit 参数
            logger.info("为 k-bit 训练准备模型 (prepare_model_for_kbit_training)...")
            # use_gradient_checkpointing 参数应与 training_args.gradient_checkpointing 一致
            model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=training_args.gradient_checkpointing)
            logger.info("模型 k-bit 准备完成。")
            # 如果使用了梯度检查点，需要确保它在 PEFT 模型中也启用
            # if training_args.gradient_checkpointing:
             #    model.gradient_checkpointing_enable() # 可能在 get_peft_model 后调用

        # 检查 lora_target_modules 是否有效
        if not peft_args.lora_target_modules:
             logger.warning("未指定 lora_target_modules，将尝试自动查找线性层 (可能不准确)。建议明确指定。")
             # 可以尝试自动查找，但这依赖于模型结构
             # modules = find_all_linear_names(model, quant_args.load_in_4bit)
             # peft_args.lora_target_modules = modules
             # logger.info(f"自动查找到的 LoRA 目标模块: {modules}")
             # raise ValueError("使用 PEFT 时必须指定 lora_target_modules。") # 或者强制要求指定

        try:
            lora_config = LoraConfig(
                r=peft_args.lora_r,
                lora_alpha=peft_args.lora_alpha,
                target_modules=peft_args.lora_target_modules,
                lora_dropout=peft_args.lora_dropout,
                bias="none", # 通常设置为 none
                task_type=TaskType.CAUSAL_LM # 指定任务类型为因果语言模型
            )
            logger.info(f"LoRA 配置: {lora_config}")

            model = get_peft_model(model, lora_config)
            logger.info("成功应用 PEFT (LoRA) 到模型。")
            model.print_trainable_parameters() # 打印可训练参数信息

            #如果使用了梯度检查点，并且 prepare_model_for_kbit_training 没有处理，在这里启用
            if training_args.gradient_checkpointing and hasattr(model, 'gradient_checkpointing_enable'):
                logger.info("为 PEFT 模型启用梯度检查点。")
                model.gradient_checkpointing_enable()


        except Exception as e:
            logger.error(f"应用 PEFT (LoRA) 失败: {e}")
            raise
    else:
        logger.info("未使用 PEFT 微调。")

    # --- 4. 初始化 Trainer ---
    # 数据整理器 (Data Collator)
    # 对于 Causal LM，使用 DataCollatorForLanguageModeling
    # mlm=False 表示不是 Masked Language Modeling
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
    logger.info("使用 DataCollatorForLanguageModeling。")

    # 初始化 Trainer
    # Trainer 会自动使用 accelerator 来处理分布式训练和设备放置

    # 准备回调列表
    callbacks = []
    if early_stopping_callback:
        callbacks.append(early_stopping_callback)
        logger.info("将 EarlyStoppingCallback 添加到 Trainer。")

    trainer = Trainer(
        model=model,                         # 可能是原始模型或 PEFT 包装后的模型
        args=training_args,                  # 官方的 TrainingArguments
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,           # 如果提供了验证集
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,                 # 传递回调列表
        # compute_metrics=... # 如果需要自定义评估指标，可以在这里提供函数
    )
    logger.info("Hugging Face Trainer 初始化完成。")

    # --- 5. 开始训练 ---
    logger.info("="*30 + " 开始训练循环 " + "="*30)
    train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
    logger.info("训练循环完成。")

    # --- 6. 保存结果 ---
    # 保存训练指标
    metrics = train_result.metrics
    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)
    logger.info(f"训练指标已保存: {metrics}")

    # 保存最终模型状态
    # Trainer 会自动处理 DeepSpeed 和 FSDP 的保存逻辑
    # 对于 PEFT，Trainer 通常只保存适配器权重
    logger.info(f"开始保存 Trainer 状态到: {training_args.output_dir}") # 日志信息修改
    trainer.save_state() # 保存 Trainer 状态 (包括优化器、调度器等)

    # --- 7. 评估 (如果需要) ---
    if training_args.do_eval and eval_dataset:
        logger.info("="*30 + " 开始评估 " + "="*30)
        eval_metrics = trainer.evaluate()
        trainer.log_metrics("eval", eval_metrics)
        trainer.save_metrics("eval", eval_metrics)
        logger.info(f"评估完成。评估指标: {eval_metrics}")

    # --- 8. 清理 ---
    # 关闭 Wandb (如果由我们手动初始化，但 Trainer 会自动处理)
    # if accelerator.is_main_process and "wandb" in training_args.report_to:
    #     wandb.finish()

    logger.info("="*30 + " 训练流程结束 " + "="*30)

    # 返回 Trainer 对象，以便主脚本可以调用 save_model
    return trainer
